# 导入必要的库和模块
import gradio as gr
import pickle
import numpy as np
from PIL import Image

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    """
    预测手写数字
    image: 从画板接收的图像
    """
    # 将图像转换为模型能理解的格式
    if image is not None:
        # 调整图像大小为8x8像素（与训练数据相同）
        image = Image.fromarray(image).convert('L')  # 转换为灰度图
        image = image.resize((8, 8))  # 调整大小
        
        # 将图像数据转换为numpy数组并归一化
        image_array = np.array(image)
        image_array = image_array.reshape(1, -1)  # 展平为1行64列
        
        # 使用模型进行预测
        prediction = knn_model.predict(image_array)
        
        return f"预测结果: {prediction[0]}"
    else:
        return "请先画一个数字！"

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
demo = gr.Interface(
    fn=predict_digit,  # 预测函数
    inputs=gr.Image(label="在手写板上画一个数字", image_mode="L", source="canvas", shape=(100, 100)),  # 输入：画板
    outputs=gr.Textbox(label="识别结果"),  # 输出：文本框显示结果
    title="手写数字识别器",
    description="用鼠标在手写板上画一个数字（0-9），AI会自动识别！",
    examples=[
        ["example_0.png"],  # 你可以添加示例图片
        ["example_1.png"]
    ]
)

# 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
    demo.launch(share=True)  # share=True 会生成一个公共链接，方便分享